<template>
  <div class="training">
    <h1>模型训练</h1>
    <div v-if="error" class="error-message">
      {{ error }}
    </div>
    <div class="training-container">
      <div class="data-panel">
        <h2>训练数据</h2>
        <div class="categories-container">
          <div v-for="(category, index) in categories" :key="index" class="category-section">
            <div class="category-header">
              <input 
                type="text" 
                v-model="category.name" 
                placeholder="输入类别名称"
                class="category-name-input"
              >
              <button 
                v-if="categories.length > 1" 
                @click="removeCategory(index)" 
                class="remove-category-btn"
              >
                删除
              </button>
            </div>
            <ImageUploader 
              :category="category.name"
              @upload="(images) => handleImagesUpload(images, index)"
              @error="handleUploadError"
            />
          </div>
        </div>
        <div class="add-category">
          <button @click="addCategory" class="add-category-btn">
            + 添加类别
          </button>
        </div>
      </div>
      <div class="config-panel">
        <h2>训练配置</h2>
        <form @submit.prevent="handleStartTraining">
          <div class="form-group">
            <label>模型选择</label>
            <select v-model="trainingConfig.model">
              <option value="mobilenet">MobileNet</option>
              <option value="resnet">ResNet</option>
              <option value="efficientnet">EfficientNet</option>
            </select>
          </div>
          <div class="form-group">
            <label>训练轮次</label>
            <input type="number" v-model="trainingConfig.epochs" min="1">
          </div>
          <div class="form-group">
            <label>学习率</label>
            <input 
              type="number" 
              v-model="trainingConfig.learningRate" 
              step="0.001" 
              min="0.001" 
              max="0.1"
            >
          </div>
          <div class="form-group">
            <label>批次大小</label>
            <input 
              type="number" 
              v-model="trainingConfig.batchSize" 
              step="8"
              min="8"
              max="128"
            >
          </div>
          <div class="form-group">
            <label>验证集比例</label>
            <input 
              type="number" 
              v-model="trainingConfig.validationSplit"
              step="0.1"
              min="0.1"
              max="0.3"
            >
          </div>
          <button type="submit" :disabled="isTraining">
            {{ isTraining ? '训练中...' : '开始训练' }}
          </button>
          <button 
            type="button" 
            @click="stopTraining" 
            v-if="isTraining"
            class="stop-btn"
          >
            停止训练
          </button>
        </form>
      </div>
      <div class="progress-panel">
        <h2>训练进度</h2>
        <div v-if="status" class="status-message">
          {{ status }}
        </div>
        <div class="progress-bar">
          <div :style="{ width: `${progress}%` }" class="progress"></div>
        </div>
        <p>{{ progress }}%</p>
        <div class="metrics-panel">
          <h3>训练指标</h3>
          <div class="metrics-chart" ref="chartRef"></div>
          <div class="current-metrics">
            <div class="metric">
              <span>当前轮次:</span>
              <span>{{ currentEpoch }}/{{ trainingConfig.epochs }}</span>
            </div>
            <div class="metric">
              <span>准确率:</span>
              <span>{{ currentAccuracy.toFixed(2) }}%</span>
            </div>
            <div class="metric">
              <span>损失值:</span>
              <span>{{ currentLoss.toFixed(4) }}</span>
            </div>
          </div>
        </div>
      </div>
      <div class="model-management-panel">
        <h2>模型管理</h2>
        <div class="model-actions">
          <div class="form-group">
            <label>模型名称</label>
            <input 
              type="text" 
              v-model="modelName" 
              placeholder="输入模型名称"
              :disabled="isTraining"
            >
          </div>
          <button 
            @click="saveCurrentModel" 
            :disabled="isTraining || !canSaveModel"
            class="save-btn"
          >
            保存模型
          </button>
        </div>

        <div class="saved-models" v-if="savedModels.length">
          <h3>已保存的模型</h3>
          <div class="models-list">
            <div v-for="model in savedModels" :key="model.name" class="model-item">
              <div class="model-info">
                <span class="model-name">{{ model.name }}</span>
                <span class="model-date">{{ model.date }}</span>
              </div>
              <div class="model-item-actions">
                <button @click="loadModel(model.name)" :disabled="isTraining">
                  加载
                </button>
                <button 
                  @click="deleteModel(model.name)" 
                  :disabled="isTraining"
                  class="delete-btn"
                >
                  删除
                </button>
              </div>
            </div>
          </div>
        </div>
      </div>
    </div>
  </div>
</template>

<script>
import ImageUploader from '../components/ImageUploader.vue'
import * as echarts from 'echarts'
import tensorflowService from '../services/tensorflow'

export default {
  name: 'TrainingView',
  components: {
    ImageUploader
  },
  data() {
    return {
      trainingConfig: {
        model: 'mobilenet',
        epochs: 10,
        learningRate: 0.001,
        batchSize: 32,
        validationSplit: 0.2
      },
      isTraining: false,
      progress: 0,
      trainingInterval: null,
      chart: null,
      currentEpoch: 0,
      currentAccuracy: 0,
      currentLoss: 1,
      trainingHistory: {
        accuracy: [],
        loss: []
      },
      images: [],
      error: null,
      categories: [
        { name: '', images: [] },
        { name: '', images: [] }
      ],
      modelName: '',
      savedModels: [],
      status: ''
    }
  },
  mounted() {
    this.initChart()
    this.loadSavedModels()
  },
  beforeUnmount() {
    if (this.chart) {
      this.chart.dispose()
    }
  },
  computed: {
    canSaveModel() {
      return this.currentEpoch > 0 && this.modelName.trim() !== '';
    }
  },
  methods: {
    initChart() {
      this.chart = echarts.init(this.$refs.chartRef)
      const option = {
        title: {
          text: '训练过程',
          left: 'center'
        },
        tooltip: {
          trigger: 'axis'
        },
        legend: {
          data: ['准确率', '损失值'],
          bottom: 0
        },
        grid: {
          left: '3%',
          right: '4%',
          bottom: '10%',
          containLabel: true
        },
        xAxis: {
          type: 'category',
          name: '轮次',
          data: []
        },
        yAxis: [
          {
            type: 'value',
            name: '准确率',
            min: 0,
            max: 100,
            position: 'left',
            axisLabel: {
              formatter: '{value}%'
            }
          },
          {
            type: 'value',
            name: '损失值',
            min: 0,
            position: 'right'
          }
        ],
        series: [
          {
            name: '准确率',
            type: 'line',
            yAxisIndex: 0,
            data: [],
            smooth: true
          },
          {
            name: '损失值',
            type: 'line',
            yAxisIndex: 1,
            data: [],
            smooth: true
          }
        ]
      }
      this.chart.setOption(option)
    },
    updateChart() {
      this.chart.setOption({
        xAxis: {
          data: Array.from({ length: this.currentEpoch + 1 }, (_, i) => i)
        },
        series: [
          {
            name: '准确率',
            data: this.trainingHistory.accuracy
          },
          {
            name: '损失值',
            data: this.trainingHistory.loss
          }
        ]
      })
    },
    async handleStartTraining() {
      try {
        // 首先检查数据是否准备好
        if (!this.preprocessData()) {
          return;
        }

        // 检查是否有足够的数据
        if (this.images.length < 2) {
          this.error = '训练数据太少，每个类别至少需要一张图片';
          return;
        }

        this.isTraining = true;
        this.error = null;
        this.progress = 0;
        this.currentEpoch = 0;
        this.currentAccuracy = 0;
        this.currentLoss = 1;
        this.trainingHistory = {
          accuracy: [],
          loss: []
        };
        
        // 准备类别信息
        const numClasses = this.categories.length;
        const categoryMap = Object.fromEntries(
          this.categories.map((cat, index) => [cat.name, index])
        );

        try {
          this.status = '正在初始化模型...';
          // 加载模型，传入类别数量
          await tensorflowService.loadModel(numClasses);
          
          this.status = '正在预处理图片数据...';
          // 预处理图片数据
          const processedData = {
            images: this.images,
            numClasses: numClasses,
            categoryMap: categoryMap
          };
          
          await tensorflowService.preprocessImages(processedData);
          
          this.status = '开始训练模型...';
          // 开始训练
          const result = await tensorflowService.train(this.trainingConfig, {
            onEpochBegin: async (epoch) => {
              this.status = `训练第 ${epoch + 1}/${this.trainingConfig.epochs} 轮...`;
            },
            onBatchEnd: async (batch) => {
              // 更新每个批次的进度
              const totalBatches = Math.ceil(this.images.length / this.trainingConfig.batchSize);
              const batchProgress = (batch + 1) / totalBatches;
              const epochProgress = this.currentEpoch / this.trainingConfig.epochs;
              this.progress = Math.round((epochProgress + batchProgress / this.trainingConfig.epochs) * 100);
            },
            onEpochEnd: async (epoch, logs) => {
              this.currentEpoch = epoch + 1;
              this.currentAccuracy = logs.acc * 100;
              this.currentLoss = logs.loss;
              
              this.trainingHistory.accuracy.push(this.currentAccuracy);
              this.trainingHistory.loss.push(this.currentLoss);
              this.updateChart();
              
              this.status = `完成第 ${epoch + 1} 轮训练`;
            }
          });
          
          this.status = '训练完成！';
          console.log('训练完成:', result);
          this.error = null;
        } catch (error) {
          console.error('训练过程出错:', error);
          this.error = error.message || '训练过程中发生错误';
          this.status = '训练失败';
          throw error;
        } finally {
          this.isTraining = false;
        }
      } catch (error) {
        console.error('训练失败:', error);
        this.error = error.message;
        this.status = '训练失败';
        this.isTraining = false;
      }
    },
    async stopTraining() {
      try {
        this.isTraining = false;
        this.progress = 0;
      } catch (error) {
        console.error('停止训练失败:', error);
        this.error = error.message;
      }
    },
    handleImagesUpload(images, categoryIndex) {
      this.categories[categoryIndex].images = images;
      this.preprocessData();
    },
    addCategory() {
      this.categories.push({ name: '', images: [] });
    },
    removeCategory(index) {
      if (this.categories.length > 1) {
        this.categories.splice(index, 1);
        this.preprocessData();
      }
    },
    preprocessData() {
      // 检查是否所有必要的数据都已准备好
      const isDataReady = this.categories.every(category => 
        category.name && category.images.length > 0
      );
      
      if (!isDataReady) {
        this.error = '请确保每个类别都有名称和图片';
        return false;
      }

      // 检查类别名称是否唯一
      const categoryNames = this.categories.map(c => c.name);
      if (new Set(categoryNames).size !== categoryNames.length) {
        this.error = '类别名称不能重复';
        return false;
      }

      // 检查每个类别是否都有足够的图片
      const hasEnoughImages = this.categories.every(category => category.images.length >= 1);
      if (!hasEnoughImages) {
        this.error = '每个类别至少需要1张图片';
        return false;
      }
      
      // 准备训练数据
      this.images = this.categories.flatMap(category => 
        category.images.map(image => ({
          ...image,
          label: category.name,
          categoryIndex: this.categories.findIndex(c => c.name === category.name)
        }))
      );

      return true;
    },
    async saveCurrentModel() {
      if (!this.modelName.trim()) {
        this.error = '请输入模型名称';
        return;
      }

      try {
        await tensorflowService.saveModel(this.modelName.trim());
        await this.loadSavedModels();
        this.error = null;
      } catch (error) {
        console.error('保存模型失败:', error);
        this.error = error.message;
      }
    },
    async loadSavedModels() {
      try {
        this.savedModels = await tensorflowService.listSavedModels();
      } catch (error) {
        console.error('获取已保存模型列表失败:', error);
        this.error = error.message;
      }
    },
    async loadModel(modelName) {
      try {
        await tensorflowService.loadModelFromIndexedDB(modelName);
        this.error = null;
        this.$message.success('模型加载成功');
      } catch (error) {
        console.error('加载模型失败:', error);
        this.error = error.message;
      }
    },
    async deleteModel(modelName) {
      try {
        await tensorflowService.deleteModel(modelName);
        await this.loadSavedModels();
        this.error = null;
      } catch (error) {
        console.error('删除模型失败:', error);
        this.error = error.message;
      }
    },
    handleUploadError(error) {
      this.error = error;
      // 3秒后清除错误信息
      setTimeout(() => {
        if (this.error === error) {
          this.error = null;
        }
      }, 3000);
    }
  }
}
</script>

<style scoped>
.training {
  padding: 20px;
}

.training-container {
  display: grid;
  grid-template-columns: 1fr 1fr 1fr;
  gap: 20px;
  margin-top: 20px;
}

.data-panel {
  background: #f5f7fa;
  padding: 20px;
  border-radius: 8px;
}

.config-panel, .progress-panel {
  padding: 20px;
  border-radius: 8px;
  background-color: #f5f5f5;
}

.form-group {
  margin-bottom: 15px;
}

.form-group label {
  display: block;
  margin-bottom: 5px;
}

.progress-bar {
  width: 100%;
  height: 20px;
  background-color: #ddd;
  border-radius: 10px;
  overflow: hidden;
}

.progress {
  height: 100%;
  background-color: #4CAF50;
  transition: width 0.3s ease;
}

.form-group input[type="number"],
.form-group select {
  width: 100%;
  padding: 8px;
  border: 1px solid #dcdfe6;
  border-radius: 4px;
  margin-top: 4px;
}

button {
  background-color: #409eff;
  color: white;
  border: none;
  padding: 10px 20px;
  border-radius: 4px;
  cursor: pointer;
  margin-right: 10px;
}

button:disabled {
  background-color: #a0cfff;
  cursor: not-allowed;
}

.stop-btn {
  background-color: #f56c6c;
}

.metrics-panel {
  margin-top: 20px;
  padding: 15px;
  background: white;
  border-radius: 4px;
  box-shadow: 0 1px 3px rgba(0,0,0,0.1);
}

.metrics-chart {
  height: 300px;
  margin: 15px 0;
}

.current-metrics {
  display: grid;
  grid-template-columns: repeat(3, 1fr);
  gap: 10px;
  text-align: center;
}

.metric {
  padding: 10px;
  background: #f5f7fa;
  border-radius: 4px;
}

.metric span:first-child {
  color: #606266;
  margin-right: 5px;
}

.metric span:last-child {
  font-weight: bold;
  color: #409eff;
}

.error-message {
  background-color: #fef0f0;
  color: #f56c6c;
  padding: 10px;
  margin-bottom: 20px;
  border-radius: 4px;
  text-align: left;
}

.categories-container {
  display: flex;
  flex-direction: column;
  gap: 20px;
  margin-bottom: 20px;
}

.category-section {
  background: white;
  border-radius: 8px;
  padding: 15px;
  box-shadow: 0 2px 8px rgba(0,0,0,0.1);
}

.category-header {
  display: flex;
  align-items: center;
  gap: 10px;
  margin-bottom: 10px;
}

.category-name-input {
  flex: 1;
  padding: 8px 12px;
  border: 1px solid #dcdfe6;
  border-radius: 4px;
  font-size: 14px;
}

.remove-category-btn {
  background: #f56c6c;
  color: white;
  border: none;
  padding: 6px 12px;
  border-radius: 4px;
  cursor: pointer;
}

.add-category {
  text-align: center;
  margin-top: 20px;
}

.add-category-btn {
  background: #409eff;
  color: white;
  border: none;
  padding: 8px 16px;
  border-radius: 4px;
  cursor: pointer;
  font-size: 14px;
}

.model-management-panel {
  padding: 20px;
  border-radius: 8px;
  background-color: #f5f5f5;
}

.model-actions {
  display: flex;
  gap: 10px;
  align-items: flex-end;
  margin-bottom: 20px;
}

.save-btn {
  background-color: #67c23a;
}

.save-btn:disabled {
  background-color: #b3e19d;
}

.models-list {
  margin-top: 10px;
}

.model-item {
  display: flex;
  justify-content: space-between;
  align-items: center;
  padding: 10px;
  background: white;
  border-radius: 4px;
  margin-bottom: 8px;
}

.model-info {
  display: flex;
  flex-direction: column;
}

.model-name {
  font-weight: bold;
  color: #303133;
}

.model-date {
  font-size: 12px;
  color: #909399;
}

.model-item-actions {
  display: flex;
  gap: 8px;
}

.model-item-actions button {
  padding: 6px 12px;
  font-size: 12px;
}

.delete-btn {
  background-color: #f56c6c;
}

.delete-btn:disabled {
  background-color: #fab6b6;
}

.status-message {
  margin: 10px 0;
  padding: 10px;
  background-color: #f0f9eb;
  color: #67c23a;
  border-radius: 4px;
  text-align: center;
}

.category-input-group {
  display: flex;
  align-items: center;
  gap: 8px;
  flex: 1;
}

.category-count {
  font-size: 14px;
  color: #909399;
  white-space: nowrap;
}
</style> 